%%% this code plots the country graphs created by make_country_graphs.m

%clear
%close all
%clc

set(0,'DefaultFigureWindowStyle','normal','DefaultFigureVisible','on') % I stopped doing the docked thing because it doesn't allow to specify size and gives inconsistent results across computers

addpath('../../Toolbox/export_fig/');
addpath('../../Toolbox/struct2csv/');

graph_width=800;
graph_height=600;

% -----------------------
% DEFINE FOLDER LOCATIONS

folders;

% save?s
save_graphs=1;

% what to plot?
basic = 0;
plot_baseline = 0;
plot_cfactuals_var = 1;
plot_slides = 0;
actual = 0;
grid = 0;
path_save_grid_plots = path_papers_graph;
osm = 0;
pop_source = 1;

% keep largest connected?
largest_connected = 1;
connected=largest_connected;

country = 'ANDES';
    
% load in the full grid mat 
graph_filename = [ path_load_grids,country,'_grid.mat' ];

load( graph_filename );

% unpack
places_grid=country_graph.places_grid;
gridmap=country_graph.gridmap;
places2=country_graph.places2;
discretized_roads=country_graph.discretized_roads;
graph_export=country_graph.graph_export;
unique_edges=country_graph.unique_edges;
edges = country_graph.edges;

%% compute some variables used in figures

discretized_avI = cell2mat( {discretized_roads.avI} )';
max_disc_avI = max( discretized_avI );

%tabulate(discretized_avI)
discretized_avI_int = round( discretized_avI );
%tabulate(discretized_avI_int)

% common parameters to all graphs
bright_line = 0.5;  % higher=brighter
bright_nodes = 0.4; % higher=brighter
max_bright = 5;    % higher=darker lines
adj_width = 0.3;
markersize = 5;
color_growth = [ 0 1 0 ];   % red = [ 1 0 0 ]  green = [ 0 1 0 ] blue = [ 0 0 1 ]
color_shrink = [ 1 0 0 ];

% number of categories for plotting
N_cat = numel( unique( cell2mat( {gridmap.quantiles} ) ) )-1;
        
adj_width = 6;
max_bright = 1.5;
GAMMA = 0.10;
CFAC_TYPE = {'mis_eng'}; %'',mis_eng 'exp_eng', 
MOBILITY = {'off'}; % 'on',
CONGESTION = {'on'};

% parameters that remain constant throughout the loop across calibrations
alpha = 0.4;
beta = 0.13;
sigma = 5;
rho = 0;
a = 1;
Ngoods = 10;
nu = 1;

for cfac_type = CFAC_TYPE

    for cong = CONGESTION

        for mobil=MOBILITY

            for jj=1:length(GAMMA)

                gamma = GAMMA(jj);

                load([path_save_cfactuals , country, '_a1_rho0_alpha0.4_sigma5_beta0.13_gamma0.1_nu1_mobil0_cong1_cfactual_', char(cfac_type), '.mat']) 
                param=cfactual.param;
                
                % load in the boundaries for each country 
                if strcmp(country, 'MERCOSUR')
                    uy =  load([ datafolder_LAC, 'UY/','country_bounds.mat']);
                    py =  load([ datafolder_LAC, 'PY/','country_bounds.mat']);
                    ar =  load([ datafolder_LAC, 'AR/','country_bounds.mat']);
                    br =  load([ datafolder_LAC, 'BR/','country_bounds.mat']);
                    br2 = shaperead([ '../../data/COUNTRY_SHP/countries_shp/countries']);
                    br = br2(strcmp({br2(:).NAME}, 'Brazil'));
                
                elseif strcmp(country, 'ANDES')
                    pe =  load(['../codes_old/country grids/Peru_grid_0_0_0.5_LAC.mat']);
                    bo =  load([ datafolder_LAC, 'BO/','country_bounds.mat']);
                    co =  load([ datafolder_LAC, 'CO/','country_bounds.mat']);
                    ec =  load([ datafolder_LAC, 'EC/','country_bounds.mat']);
                end
                
                % recover results and g
                g = cfactual.g;
                results_actual = cfactual.results_actual;
                results_cfactual = cfactual.results_cfactual;

                % GDP per capita in traded sector in calibrated model, cfactual, and data
                y_calib = sum( results_actual.Pjn.*results_actual.Yjn,2 )./results_actual.Lj;
                y_exp = sum( results_cfactual.Pjn.*results_cfactual.Yjn,2 )./results_cfactual.Lj;
                y_data = g.Y./g.L;

                % total GDP in calibration and cfactual
                PHj_calib = (1-param.alpha) / param.alpha * results_actual.PCj .* results_actual.cj ./ cfactual.param.hj;
                GDP_calib = sum(results_actual.Pjn.*results_actual.Yjn,2) + PHj_calib.*cfactual.param.Hj;

                PHj_exp = (1-param.alpha) / param.alpha * results_cfactual.PCj .* results_cfactual.cj ./ cfactual.param.hj;
                GDP_exp = sum(results_cfactual.Pjn.*results_cfactual.Yjn,2) + PHj_exp.*cfactual.param.Hj;

                % population in calibrated model, cfactual, and data
                L_calib = results_actual.Lj;
                L_exp = results_cfactual.Lj;
                L_data = g.L;

                % total income in calibrated model, cfactual, and data
                Y_calib = L_calib.*y_calib;
                Y_exp = L_exp.*y_exp;
                Y_data = L_data.*y_data;

                % total consumption in calibrated and cfactual
                C_calib = results_actual.Cj;
                C_exp = results_cfactual.Cj;

                % per capita
                c_calib = C_calib./L_calib;
                c_exp = C_exp./L_exp;

                % welfare gain
                welfare_gain = consumption_equivalent(cfactual.param,g,c_calib,L_calib,results_cfactual.welfare)-1;

                % exports and imports in calib
                [ X_calib,M_calib ] = recover_X_M( results_actual,g );

                % exports and imports in cfactual
                [ X_exp,M_exp ] = recover_X_M( results_actual,g );

                % X/GDP
                XY_calib = X_calib./Y_calib;
                XY_exp = X_exp./Y_exp;

                % M/GDP
                MY_calib = M_calib./Y_calib;
                MY_exp = M_exp./Y_exp;

                % net exports over GDP
                nx_calib = ( X_calib-M_calib )./Y_calib;
                nx_exp = ( X_exp-M_exp )./Y_exp;

                % fundamentals
                H_calib = cfactual.param.Hj;
                Z_calib = max( cfactual.param.Zjn,[],2 );

                % infrastructure long
                avI_obs = triu(g.avI);
                avI_obs = avI_obs(:);

                avI_exp = triu(results_cfactual.Ijk);
                avI_exp = avI_exp(:);

                % quantiles in calibrated and counterfactual model
                quantile_L_calib = assign_quantiles( L_calib );
                quantile_L_exp = assign_quantiles( L_exp );

                % lanes in calibrated and counterfactual model
                I_obs = zeros( length(unique_edges),1 );
                I_exp = zeros( length(unique_edges),1 );
                for j=1:length(unique_edges)

                    I_obs(j) = g.avI( unique_edges(j,1),unique_edges(j,2) );
                    I_exp(j) = results_cfactual.Ijk( unique_edges(j,1),unique_edges(j,2) );


                end

                % identify producers of differentiated goods
                diff_producers = find( results_actual.Yjn(:,1)==0 );  % indexes of differentiated producers

                % check dI adds up to 0.5
                % total_dI = sum( sum( g.delta_i.*( results_exp.I-g.avI ) ) )

                % all these cases have no construction in any case
                %I_obs( I_obs==0.0001 )=0;
                %I_exp( I_obs==0 )=0;

                % set maximum expansion to 10
                %maxI_obs = max( max( I_obs ),10 );
                %I_exp = min( I_exp,maxI_obs );  % bound counterfactual I

                % compute expansion
                gL = ( L_exp-L_calib )./( 0.5*( L_exp+L_calib ) );
                gc = ( c_exp-c_calib )./( 0.5*( c_exp+c_calib ) );
                gy = ( y_exp-y_calib )./( 0.5*( y_exp+y_calib ) );
                gXY = ( XY_exp-XY_calib )./( 0.5*( XY_exp+XY_calib ) );
                gMY = ( MY_exp-MY_calib )./( 0.5*( MY_exp+MY_calib ) );
                gnx = ( nx_exp-nx_calib )./( 0.5*( nx_exp+nx_calib ) );
                dI = I_exp-I_obs;
                dL =  L_exp-L_calib;
                gL2 = ( L_exp-L_calib )./( L_calib  );
                gI = (I_exp-I_obs)./((I_exp+I_obs)*0.5);

                if strcmp(cfac_type, 'mis_eng') && ~strcmp(country, 'Mexico')
                    dI(dI>prctile(dI, 90)) = prctile(dI, 90);
                    dI(dI<prctile(dI, 10)) = prctile(dI, 10);

                elseif strcmp(cfac_type, 'mis_eng') && strcmp(country, 'Mexico')
                    dI(dI>prctile(dI, 90)) = prctile(dI, 90);
                    dI(dI<prctile(dI, 5)) = prctile(dI, 5);
                end

                if strcmp(cfac_type, 'exp_eng')
                    dI(dI>prctile(dI, 90)) = prctile(dI, 90);
                    % dI(dI<prctile(dI, 10)) = prctile(dI, 10);
                end

                % max I
                maxI_obs = max( max(I_obs),max(I_exp) );
                %[ I_exp( I_obs==0.0001 ) I_obs( I_obs==0.0001 ) I_exp( I_obs==0.0001 )./I_obs( I_obs==0.0001 ) ]

                show_cells = 0;  % otherwise, show nodes

                %% plot change in network and in population

                close(figure(6))
                h=figure(6);
                set(h,'Position',[0 0 graph_width graph_height]);

                if strcmp(country, 'ANDES')
                    mapshow( pe.country_graph.country_bounds,...
                                'FaceColor','black' )
                end

                mapshow( gridmap,...
                    'FaceColor','Black',...
                    'EdgeColor','White','LineWidth', 0.5 )
                
                if strcmp(country, 'MERCOSUR')
                    mapshow( uy.country_bounds,...
                        'FaceColor','none',...
                        'EdgeColor','m','LineWidth',  2.5 )

                    mapshow( py.country_bounds,...
                        'FaceColor','none',...
                        'EdgeColor','m','LineWidth',  2.5 )

                    mapshow( ar.country_bounds,...
                        'FaceColor','none',...
                        'EdgeColor','m','LineWidth',  2.5 )

                    mapshow( br,...
                        'FaceColor','none',...
                        'EdgeColor','m','LineWidth', 2.5 )
                elseif strcmp(country, 'ANDES')
                    
                    mapshow( pe.country_graph.country_bounds,...
                    'FaceColor','none',...
                    'EdgeColor','m','LineWidth',  2.5 )

                    mapshow( bo.country_bounds,...
                                    'FaceColor','none',...
                                    'EdgeColor','m','LineWidth',  2.5 )

                    mapshow( co.country_bounds,...
                                    'FaceColor','none',...
                                    'EdgeColor','m','LineWidth',  2.5 )

                    mapshow( ec.country_bounds,...
                                    'FaceColor','none',...
                                    'EdgeColor','m','LineWidth',  2.5 )    
                end
                
                
                % population
                x = dL;
                x(x<prctile(x,3.5)) = prctile(x,3.5);
                x(x>prctile(x,96.5)) = prctile(x,96.5);
                if show_cells && param.mobility
                    for i=1:length( gridmap )

                        norm_neg = 1/abs( mean(x(x<0)) )*1/2.5;   %gL(gL<0)*norm_neg
                        norm_pos = 1/max( x(x>0) )*2.5;         %gL(gL>0)*norm_neg

                        relsize = x(i);
                        bring_fire = 0;

                        color = min( max( -relsize*norm_neg,0 )*color_shrink+...
                            ( max( -relsize*norm_neg,0.5 )-0.5 )*bring_fire*[ 0 1 0 ],1 )+...
                            min( max( relsize*norm_pos,0 )*color_growth+...
                            ( max( relsize*norm_pos,0.5 )-0.5 )*bring_fire*[ 0 1 0 ],1 );

                        mapshow( gridmap( i ),...
                            'FaceColor',color );

                    end
                end

                % lanes
                for j=1:length( dI )
                    [ j ]
                    % COLOR SCALE OF LANES
%                             colorweight = abs( dI(j) )/max_bright;
                    colorweight = abs( dI(j) )/max(abs(dI))/max_bright;
                    lanewidth = abs( dI(j) )/max(abs(dI))*adj_width;

                    if dI(j)>0
                        color =  min( bright_line+( 1-bright_line )*colorweight*color_growth,1 );
                        lane_size = lanewidth;
                    elseif dI(j)<0
                        color =  min( bright_line+( 1-bright_line )*3*colorweight*color_shrink,1 );
                        lane_size = lanewidth;
                    end

                      mapshow( discretized_roads( j ),...
                                    'Color',min( color,1 ),...
                                    'LineWidth',lane_size )
                                    
                    %end
                    hold on;

                end


                if ~show_cells
                    for i=1:length( gridmap )
                        [i ]
                        norm_neg = 1/abs( mean(x(x<0)) )*1/2.5;   %gL(gL<0)*norm_neg
                        norm_pos = 1/max( x(x>0) )*2.5;         %gL(gL>0)*norm_neg

                        relsize = x(i);
                        bring_fire = 0;

                        if param.mobility
                            color = min( max( -relsize*norm_neg,0 )*color_shrink+...
                                ( max( -relsize*norm_neg,0.5 )-0.5 )*bring_fire*[ 0 1 0 ],1 )+...
                                min( max( relsize*norm_pos,0 )*color_growth+...
                                ( max( relsize*norm_pos,0.5 )-0.5 )*bring_fire*[ 0 1 0 ],1 );
                        else
                            color = [ 0 0 0 ];
                        end

                        mapshow( places2( i ),...
                        'Marker','o',...
                        'MarkerFaceColor',color,...
                        'MarkerEdgeColor',[ 1 1 1 ],...
                        'MarkerSize',markersize )

                    end
                end

                margin = 0.5;
                Xlim = [ min( cell2mat({gridmap.X}) )-margin;max( cell2mat({gridmap.X}) )+margin ];
                Ylim = [ min( cell2mat({gridmap.Y}) )-margin;max( cell2mat({gridmap.Y}) )+margin ];
                set( gca,'XTick',[],'YTick',[],'XLim',Xlim,'YLim',Ylim,'Box','on' )

                if save_graphs
                   % export_fig([ path_papers_graph,filename,'_cfactual_dL_dI_',char(cfac_type),'.eps' ], '-depsc');
                     export_fig([ '../../../notes/images/',country,'_',char(cfac_type),'.eps' ], '-depsc');
                end
            end

        end

    end

end



